{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting Started with Medical Imaging\n",
    "\n",
    "In this tutorial, we will be looking at how to work with Medical Imaging samples in FiftyOne. We will specifically looking at DICOM files and CT scan files! While there are a few known variations on how they can be store, the pattern is the same.\n",
    "\n",
    "Note, we will not be covering how to load medical images that are normally stored as `.jpg` or `.png` . Please refer to the quickstart if you are working with normal images!\n",
    "\n",
    "## Installation\n",
    "Some new libraries are required to load the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pydicom==2.4.4 rt_utils kagglehub nibabel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Downloading DICOM Dataset\n",
    "\n",
    "First let's grab some example DICOM files. Below is a [great demo dataset](https://www.kaggle.com/datasets/aryashah2k/hippocampal-sparing-dataset) from Kaggle. It is not of the highest quality, but it works well for beginners. Inside it, we have 25 brain scans of paitents with annotations of their left and right hippocampus stored in `.dcm` files. Lets begin by downloading the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import kagglehub\n",
    "\n",
    "# Download latest version\n",
    "path = kagglehub.dataset_download(\"aryashah2k/hippocampal-sparing-dataset\")\n",
    "\n",
    "print(\"Path to dataset files:\", path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading DICOM Files\n",
    "\n",
    "FiftyOne features a [DICOM dataset loader](https://docs.voxel51.com/user_guide/dataset_creation/datasets.html#dicomdataset) that we will be taking advantage of here. This will load our `.dcm` files in as 2D slices that we can then play as a video once its all loaded. \n",
    "\n",
    "An important note here is the use of `dicom_path`. In our dataset we have MR dicoms and RS dicoms, with the former being the scan and the latter being the annotations. We only want to load the scan as our base sample and load the annotations on later. FiftyOne is very flexible with how you load and present data, so feel free to dive more into the [dataset loader](https://docs.voxel51.com/user_guide/dataset_creation/datasets.html#dicomdataset) to get the best experience for you!\n",
    "\n",
    "Let's begin loading a single patient's data below. Check out the comments to find out what is happening at each step!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import fiftyone as fo\n",
    "import glob\n",
    "from rt_utils import RTStructBuilder\n",
    "import numpy as np\n",
    "import cv2\n",
    "import os\n",
    "\n",
    "\n",
    "# Path to the dataset\n",
    "dataset_dir = path + \"/HippocampalMRISlices/01/\"\n",
    "\n",
    "# Create the dataset\n",
    "dataset = fo.Dataset.from_dir(\n",
    "    dataset_dir=dataset_dir,\n",
    "    dataset_type=fo.types.DICOMDataset,\n",
    "    dicom_path=\"MR*.dcm\",\n",
    "    name=\"Patient1Scans\",\n",
    ")\n",
    "\n",
    "# Load RTStruct for labels\n",
    "rtstruct_path = glob.glob(dataset_dir + \"RS*.dcm\", recursive=True)\n",
    "rtstruct = RTStructBuilder.create_from(dicom_series_path=dataset_dir, rt_struct_path=rtstruct_path[0])\n",
    "\n",
    "# Get structure names\n",
    "structures = rtstruct.get_roi_names()\n",
    "print(\"Structures in RTStruct:\", structures)\n",
    "\n",
    "# Get mask for each structure\n",
    "L_mask3d = rtstruct.get_roi_mask_by_name(structures[0])\n",
    "R_mask3d = rtstruct.get_roi_mask_by_name(structures[1])\n",
    "\n",
    "# Sort dataset from bottom to top of the scan\n",
    "view = dataset.sort_by(\"SliceLocation\")\n",
    "\n",
    "# Add masks to the dataset\n",
    "i = 0\n",
    "for sample in view:\n",
    "\n",
    "    # Load left Hippocampus mask\n",
    "    sample[\"Hippocampus_L\"] = fo.Segmentation(mask=L_mask3d[:,:,i].astype(np.uint8))\n",
    "\n",
    "    # Load right Hippocampus mask\n",
    "    sample[\"Hippocampus_R\"] = fo.Segmentation(mask=R_mask3d[:,:,i].astype(np.uint8))\n",
    "    i = i + 1\n",
    "    sample.save()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Afterwards we can open up our app to see patient 1's scan and segmentation! Right away you might notice that we are only looking at individual slices. If we want to watch the playback of these, we need to use dynamic grouping! \n",
    "\n",
    "To dynamicly group our dataset, click the dynamic group button above the grid that looks like two merging arrows. We want to group by `Patient ID` and order by `SliceLocation`. This should present now just one sample in the grid. Lastly, click the gear icon and check \"Render frames as video\". Below is a video showing how to do the above as well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session = fo.launch_app(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![patient1_slices](https://cdn.voxel51.com/patient1_slices.webp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Automatic Metadata Ingestion\n",
    "\n",
    "One of the other awesome features of the DICOM Dataset Loader is that it will automatically grab all the meta data from our `.dcm` file for us! Scan the sidebar in the app to search, filter and sort across fields like `BodyPartExamined`, `Manufacturer`, `SliceLocation` and more!\n",
    "\n",
    "But to get a full picture, we should load in the other 24 patients! To do this we will follow the same pattern above we did for Patient 1 and the use [merge_sample](https://docs.voxel51.com/api/fiftyone.core.dataset.html#fiftyone.core.dataset.Dataset.merge_sample) in order to combine datasets into one final dataset for us!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get a list of all directories in the dataset\n",
    "directories = [d for d in os.listdir(path + \"/HippocampalMRISlices/\") if os.path.isdir(os.path.join(path + \"/HippocampalMRISlices/\", d))]\n",
    "\n",
    "# Create a new dataset to merge all the datasets\n",
    "final_dataset = fo.Dataset(\"Hippocampal Contouring Dataset 2\")\n",
    "\n",
    "# Loop through all directories\n",
    "for directory in directories:\n",
    "    dataset_dir = path + \"/HippocampalMRISlices/\" + directory + \"/\"\n",
    "\n",
    "    # Create the dataset\n",
    "    dataset = fo.Dataset.from_dir(\n",
    "        dataset_dir=dataset_dir,\n",
    "        dataset_type=fo.types.DICOMDataset,\n",
    "        dicom_path=\"MR*.dcm\"\n",
    "    )\n",
    "\n",
    "    # Load RTStruct for labels\n",
    "    rtstruct_path = glob.glob(dataset_dir + \"RS*.dcm\", recursive=True)\n",
    "    rtstruct = RTStructBuilder.create_from(dicom_series_path=dataset_dir, rt_struct_path=rtstruct_path[0])\n",
    "\n",
    "    # Get structure names\n",
    "    structures = rtstruct.get_roi_names()\n",
    "\n",
    "    # Get mask for each structure\n",
    "    L_mask3d = rtstruct.get_roi_mask_by_name(structures[0])\n",
    "    R_mask3d = rtstruct.get_roi_mask_by_name(structures[1])\n",
    "\n",
    "    # Sort dataset from bottom to top of the scan\n",
    "    view = dataset.sort_by(\"SliceLocation\")\n",
    "\n",
    "    i = 0\n",
    "    for sample in view:\n",
    "\n",
    "        # Load left Hippocampus mask\n",
    "        sample[\"Hippocampus_L\"] = fo.Segmentation(mask=L_mask3d[:,:,i].astype(np.uint8))\n",
    "\n",
    "        # Load right Hippocampus mask\n",
    "        sample[\"Hippocampus_R\"] = fo.Segmentation(mask=R_mask3d[:,:,i].astype(np.uint8))\n",
    "        i = i + 1\n",
    "        sample.save()\n",
    "    \n",
    "    # Merge the dataset with the final dataset\n",
    "    final_dataset.merge_samples(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have all the patients into a single dataset, we can change our app to view our current dataset! Try dynamic grouping to see all of our different scans in our dataset as videos!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.dataset = final_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![all_slices](https://cdn.voxel51.com/all_slices.webp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with CT Scans\n",
    "In FiftyOne, working with CT Scans can be a bit more tricky. CT Scans can come in multiple different file formats that all have slight variations. In order to ingest these files, we need to manually slice them ourselves. The main file formats used for CT scans are DICOM (`.dcm`), NIfTI (`.nii` or `nii.gz`), NRRD (`.nrrd`), or occasionally MHA/MHD (`.mha`, `.mhd` + `.raw`). \n",
    "\n",
    "The good new is if it is stored in DICOM format, you can follow the instructions above! If it is one of the other formats, we need to define what are base atomic sample is and store that as image. This means loading a NIfTI file for example, defining the slice in python, and then saving it back with `opencv` or a similar tool. We will show an example below using NIfTI files with masks.\n",
    "\n",
    "Note, there are ongoing community efforts to bring in 3D visualizers for Medical Imaging to better support DICOM, NIfTI, NRRD, and more! If you are interested in testing out experimental tools or contributing yourself, checkout the work at our [Discord](https://community.voxel51.com)!\n",
    "\n",
    "## Downloading the CT Scan Data\n",
    "We will be using another [demo dataset from kaggle](https://www.kaggle.com/datasets/andrewmvd/covid19-ct-scans?select=ct_scans) that captures some COVID-19 lung scans! Download using the code below:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import kagglehub\n",
    "\n",
    "# Download latest version\n",
    "path = kagglehub.dataset_download(\"andrewmvd/covid19-ct-scans\")\n",
    "\n",
    "print(\"Path to dataset files:\", path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now with our data downloaded, just like before lets first begin with just a single patient scan."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nibabel as nib\n",
    "import numpy as np\n",
    "import os\n",
    "import cv2\n",
    "import fiftyone as fo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading NIfTI Files in FiftyOne\n",
    "We are going to define some helper functions to handle to loading of the NIfTI files. This time around, we will take a different approach to working with medical images. Instead of saving each individual slice like before, we will instead save an `mp4` of all the slices. This will be memory efficient and easy for us to slice and search through.\n",
    "\n",
    "In our first helper function, `nii_to_mp4`, we will take in a `.nii` file from our CT scan and save it back as an mp4 in the axial view. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def nii_to_mp4(filepath, fps=30):\n",
    "    '''\n",
    "    Reads .nii file and returns pixel array\n",
    "    '''\n",
    "    # Load .nii file\n",
    "    ct_scan = nib.load(filepath)\n",
    "    array   = ct_scan.get_fdata()\n",
    "\n",
    "    # Rotate array\n",
    "    array   = np.rot90(np.array(array))\n",
    "\n",
    "    # Normalize pixel values\n",
    "    array = (array - np.min(array)) / (np.max(array) - np.min(array)) * 255\n",
    "    array = array.astype(np.uint8)\n",
    "\n",
    "    # Define video writer\n",
    "    height, width, depth = array.shape\n",
    "    output_path = filepath.replace(\".nii\", \".mp4\")\n",
    "    fourcc = cv2.VideoWriter_fourcc(*'mp4v')\n",
    "    video_writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height), isColor=True)\n",
    "\n",
    "    # Write each axial slice to the video\n",
    "    for i in range(depth):\n",
    "        frame = array[:, :, i]\n",
    "        # Convert grayscale to BGR for consistent writing\n",
    "        frame_bgr = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)\n",
    "        video_writer.write(frame_bgr)\n",
    "\n",
    "    video_writer.release()\n",
    "    return(output_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In our next helper function, `read_nii`, we load in a `.nii` file like before, but this time instead we are going to return the array to be loaded onto our samples as a Segmentation mask!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_nii(filepath):\n",
    "    '''\n",
    "    Reads .nii file and returns pixel array\n",
    "    '''\n",
    "    ct_scan = nib.load(filepath)\n",
    "    array   = ct_scan.get_fdata()\n",
    "    array   = np.rot90(np.array(array))\n",
    "    array = array.astype(np.uint8)  \n",
    "    return(array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With our helper functions, we can grab our first sample and attempt to load it in! Dont forget we need the 3 different segmentation masks as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert Scan to mp4\n",
    "sample_ct   = nii_to_mp4(path + \"/ct_scans/coronacases_org_001.nii\")\n",
    "\n",
    "# Read the masks\n",
    "sample_lung = read_nii(path + \"/lung_mask/coronacases_001.nii\")\n",
    "sample_infe = read_nii(path + \"/infection_mask/coronacases_001.nii\")\n",
    "sample_all  = read_nii(path + \"/lung_and_infection_mask/coronacases_001.nii\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we create a FiftyOne sample. Since we passed in a `.mp4` file, this will be a video sample. With video samples, we can store any type of label on each frame. Let's try loading our segmentation masks on each slice of our CT scan!\n",
    "\n",
    "If you want to learn more about how to work with FiftyOne video samples and datasets, check out [Getting Started with Videos](https://voxel51.com/webinars)!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create output folder\n",
    "sample = fo.Sample(filepath=sample_ct)\n",
    "\n",
    "\n",
    "for i in range(sample_lung.shape[2]):\n",
    "\n",
    "    # Add masks to each frame. Note that video frames start from 1, there is no frame 0!\n",
    "    sample.frames[i+1][\"Infection Mask\"] = fo.Segmentation(mask=sample_infe[:, :, i].astype(np.uint8))\n",
    "    sample.frames[i+1][\"Lung Mask\"] = fo.Segmentation(mask=sample_lung[:, :, i].astype(np.uint8))\n",
    "    sample.frames[i+1][\"All Masks\"] = fo.Segmentation(mask=sample_all[:, :, i].astype(np.uint8))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With our video sample created and annotations added, we can create a dataset to load into. The final step is to make sure that the `.mp4` we saved can be rendered in all browsers. Not all `mp4`s are encoded the same, so we use `fo.utils.video.reencode_videos()` to ensure that it will render for us in the app! Speaking of which, after it is reencoded, lets open up the app and look at our first scan!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a dataset of Patient 1\n",
    "dataset = fo.Dataset(name=\"COVID-19 CT Scans Patient 1\", overwrite=True)\n",
    "dataset.add_samples([sample])\n",
    "\n",
    "# Reencode videos so it will play in any browser\n",
    "fo.utils.video.reencode_videos(dataset)\n",
    "\n",
    "# Launch the App\n",
    "session = fo.launch_app(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![patient1_ct1](https://cdn.voxel51.com/patient1_ct.webp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great! Now we can visualize our first CT Scan. Try different color options by clicking the color palatte button above the grid and change \"Color annotations by\" to `value`.\n",
    "\n",
    "Next step is to repeat the same process for our annotations as well and add them to our sample!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List all files in the directory that end with .nii\n",
    "ct_scan_files = [f for f in os.listdir(path + \"/ct_scans/\") if f.endswith(\".nii\") and os.path.isfile(os.path.join(path + \"/ct_scans/\", f))]\n",
    "ct_scan_files.sort()\n",
    "print(ct_scan_files[0])\n",
    "\n",
    "lung_files = [f for f in os.listdir(path + \"/lung_mask/\") if f.endswith(\".nii\") and os.path.isfile(os.path.join(path + \"/lung_mask/\", f))]\n",
    "lung_files.sort()\n",
    "print(lung_files[0])\n",
    "\n",
    "infe_files = [f for f in os.listdir(path + \"/infection_mask/\") if f.endswith(\".nii\") and os.path.isfile(os.path.join(path + \"/infection_mask/\", f))]\n",
    "infe_files.sort()\n",
    "print(infe_files[0])\n",
    "\n",
    "all_files = [f for f in os.listdir(path + \"/lung_and_infection_mask/\") if f.endswith(\".nii\") and os.path.isfile(os.path.join(path + \"/lung_and_infection_mask/\", f))]\n",
    "all_files.sort()\n",
    "print(all_files[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have all of our scans sorted and in a list, it is now time to loop over them and create some samples! \n",
    "\n",
    "This pattern of creating samples, appending them to a list, and the adding them to a FiftyOne dataset is one of the fastest ways to create a FiftyOne dataset in terms of time complexity. This is because it minimizes the number of writes to your backend database. Try using it anytime you are loading large number of samples to cut down the load times!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = []\n",
    "for i in range(len(ct_scan_files)):\n",
    "    # Convert Scan to mp4\n",
    "    sample_ct   = nii_to_mp4(path + \"/ct_scans/\" + ct_scan_files[i])\n",
    "\n",
    "    # Read the masks\n",
    "    sample_lung = read_nii(path + \"/lung_mask/\" + lung_files[i])\n",
    "    sample_infe = read_nii(path + \"/infection_mask/\" + infe_files[i])\n",
    "    sample_all  = read_nii(path + \"/lung_and_infection_mask/\" + all_files[i])\n",
    "\n",
    "    # Create Sample\n",
    "    sample = fo.Sample(filepath=sample_ct)\n",
    "\n",
    "    for i in range(sample_lung.shape[2]):\n",
    "\n",
    "        # Add masks to each frame\n",
    "        sample.frames[i+1][\"Infection Mask\"] = fo.Segmentation(mask=sample_infe[:, :, i].astype(np.uint8))\n",
    "        sample.frames[i+1][\"Lung Mask\"] = fo.Segmentation(mask=sample_lung[:, :, i].astype(np.uint8))\n",
    "        sample.frames[i+1][\"All Masks\"] = fo.Segmentation(mask=sample_all[:, :, i].astype(np.uint8))\n",
    "    \n",
    "    samples.append(sample)\n",
    "\n",
    "# Create a dataset of all Patients\n",
    "dataset = fo.Dataset(name=\"COVID-19 CT Scans\", overwrite=True)\n",
    "dataset.add_samples(samples)\n",
    "\n",
    "# Reencode videos so it will play in any browser\n",
    "fo.utils.video.reencode_videos(dataset)\n",
    "\n",
    "# Launch the App\n",
    "session = fo.launch_app(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After running the above code, you should be able to visualize all the scans in the dataset now!\n",
    "\n",
    "![all_scans](https://cdn.voxel51.com/all_patients_ct.webp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next up, try checking out some of our other docs to take your Medical Imaging dataset to next level by utilitzing tools like:\n",
    "- Visualizing Embeddings\n",
    "- Similarity Search\n",
    "- MedSAM2 in Model Zoo"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "OSS310",
   "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.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}