multiple-labels.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0fb35ca1-7e8a-4527-8d0c-ad94995de5c0",
   "metadata": {},
   "source": [
    "# Ultrack I2K 2023 - Multiple hypotheses tracking\n",
    "\n",
    "This tutorial shows the multiple hypotheses tracking capabilities of Ultrack. \n",
    "\n",
    "Here, rather than searching for an optimal segmentation parameter, we sampled multiple segmentations with different parametrizations and used Ultrack to find the best segments, obtaining more accurate cell tracking.\n",
    "\n",
    "\n",
    "## Setting up Colab runtime\n",
    "\n",
    "If you are using Colab, we recommend to set up the runtime to use a GPU.\n",
    "To do so, go to `Runtime > Change runtime type` and select `GPU` as the hardware accelerator.\n",
    "\n",
    "## Setup Dependencies\n",
    "\n",
    "This step is only necessary if you are on Colab or don't have the required packages.\n",
    "\n",
    "IMPORTANT: The runtime must be initialized.\n",
    "\n",
    "Uncomment and run the following commands to install all required packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "017c2766-fc33-42f9-a0f2-30fda422c59f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install stackview cellpose 'napari[all]' ultrack ipycanvas==0.11 cucim\n",
    "# !pip install git+https://github.com/Janelia-Trackathon-2023/traccuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37cbc363-365c-4b9c-8816-a7858ef4603c",
   "metadata": {},
   "source": [
    "## Download Dataset\n",
    "\n",
    "Download the Fluo-C2DL-Huh7 dataset from the [Cell Tracking Challenge](celltrackingchallenge.net), which contains fluorescence microscopy images for cell tracking.\n",
    "\n",
    "The dataset will be used for demonstrating the segmentation and tracking workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5ff4739",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -nc http://data.celltrackingchallenge.net/training-datasets/Fluo-C2DL-Huh7.zip\n",
    "!unzip -n Fluo-C2DL-Huh7.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ea37619-cca3-4029-8189-456fc9153fcf",
   "metadata": {},
   "source": [
    "## Import Libraries\n",
    "\n",
    "Import the libraries needed for reading images, processing them, cell segmentation, tracking, and performance metrics. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "923c68b7",
   "metadata": {
    "id": "923c68b7"
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from typing import Dict\n",
    "from cellpose.models import Cellpose\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import stackview\n",
    "from dask.array.image import imread\n",
    "from numpy.typing import ArrayLike\n",
    "from rich import print\n",
    "\n",
    "from traccuracy import run_metrics\n",
    "from traccuracy.loaders import load_ctc_data\n",
    "from traccuracy.matchers import CTCMatched\n",
    "from traccuracy.metrics import CTCMetrics\n",
    "\n",
    "from ultrack import track, to_tracks_layer, tracks_to_zarr, to_ctc\n",
    "from ultrack.utils import labels_to_edges\n",
    "from ultrack.config import MainConfig\n",
    "from ultrack.imgproc import normalize\n",
    "from ultrack.imgproc.segmentation import Cellpose\n",
    "from ultrack.utils.array import array_apply"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7905b69-db4b-4b45-9451-5aa3248a42d9",
   "metadata": {},
   "source": [
    "## Colab or Local\n",
    "\n",
    "Change the `COLAB` variable to `True` or `False` depending on whether you are running this notebook on Colab or locally.\n",
    "\n",
    "When running locally napari will be used a the image viewer, while on Colab the images will be displayed using `stackview`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63901b1c-ac75-4873-8ce2-757da85bd754",
   "metadata": {},
   "outputs": [],
   "source": [
    "COLAB = True\n",
    "# COLAB = False\n",
    "\n",
    "if COLAB:\n",
    "    viewer = None\n",
    "\n",
    "    # fixes colab encoding error\n",
    "    import locale\n",
    "    locale.getpreferredencoding = lambda: \"UTF-8\"\n",
    "\n",
    "    # enabling colab output\n",
    "    try:\n",
    "        from google.colab import output\n",
    "        output.enable_custom_widget_manager()\n",
    "    except ModuleNotFoundError as e:\n",
    "        print(e)\n",
    "else:\n",
    "    import napari\n",
    "    from napari.utils import nbscreenshot\n",
    "\n",
    "    viewer = napari.Viewer()\n",
    "\n",
    "    def screenshot() -> None:\n",
    "        display(nbscreenshot(viewer))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b381ca65-f913-4ac3-a23b-66faa532fc54",
   "metadata": {},
   "source": [
    "## Load Data\n",
    "\n",
    "Load the Fluo-C2DL-Huh7 dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "291788ce",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "291788ce",
    "outputId": "e6f14189-e280-49cb-b912-e47cd860c985"
   },
   "outputs": [],
   "source": [
    "dataset = \"02\"\n",
    "path = Path(\"Fluo-C2DL-Huh7\") / dataset\n",
    "image = imread(str(path / \"*.tif\"))\n",
    "\n",
    "if COLAB:\n",
    "    display(stackview.slice(image))\n",
    "else:\n",
    "    viewer.add_image(image)\n",
    "    screenshot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e189858e-55e2-42e8-9a4d-94c511d347fe",
   "metadata": {},
   "source": [
    "## Configuration\n",
    "\n",
    "We'll use the same configuration as in the previous example, except for `config.segmentation_config.min_frontier` which had its value decreased.\n",
    "\n",
    "The `min_frontier` merges regions with an average contour lower than the provided value.\n",
    "Since the contours are combined by averaging, the previous value of 0.1 removed relevant segments from the candidate hypotheses.\n",
    "\n",
    "As a reminder, the configuration parameters documentation can be found [here](https://github.com/royerlab/ultrack/blob/main/ultrack/config/README.md)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff451bcc",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 668
    },
    "id": "ff451bcc",
    "outputId": "0cb3ebe4-e371-4b19-9d67-25e35e7f8b86"
   },
   "outputs": [],
   "source": [
    "config = MainConfig()\n",
    "\n",
    "# Candidate segmentation parameters\n",
    "config.segmentation_config.n_workers = 8\n",
    "config.segmentation_config.min_area = 2500\n",
    "config.segmentation_config.min_frontier = 0.05  # NOTE: this parameter is not the same as in intro.ipynb\n",
    "\n",
    "# Setting the maximum number of candidate neighbors and maximum spatial distance between cells\n",
    "config.linking_config.max_neighbors = 5\n",
    "config.linking_config.max_distance = 100\n",
    "config.linking_config.n_workers = 8\n",
    "\n",
    "# Adding absurd weight to division because there's no diving cell\n",
    "config.tracking_config.division_weight = -100\n",
    "# Very few tracks enter/leave the field of view, increasing penalization\n",
    "config.tracking_config.disappear_weight = -1\n",
    "config.tracking_config.appear_weight = -1\n",
    "\n",
    "print(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "223e0559-88ac-480d-8585-235d7b3d31c9",
   "metadata": {},
   "source": [
    "## Cellpose Segmentation\n",
    "\n",
    "The same function as `intro.ipynb` to segment cells within each frame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7829b39a-5471-4eb8-a8d9-2ce704445a75",
   "metadata": {},
   "outputs": [],
   "source": [
    "cellpose = Cellpose(model_type=\"cyto2\", gpu=True)\n",
    "\n",
    "def predict(frame: ArrayLike, gamma: float) -> ArrayLike:\n",
    "    norm_frame = normalize(np.asarray(frame), gamma=gamma)\n",
    "    return cellpose(norm_frame, tile=False, normalize=False, diameter=75.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04c3717e-7d55-4695-a15d-5443ebd27ac0",
   "metadata": {},
   "source": [
    "## Metrics\n",
    "\n",
    "Helper function to evaluate tracking score using [Cell Tracking Challenge](celltrackingchallenge.net)'s metrics and annotations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa88e34c-5787-48df-9adf-9f2e2f4e74e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "gt_path = path.parent / f\"{dataset}_GT\"\n",
    "gt_data = load_ctc_data(gt_path / \"TRA\")\n",
    "\n",
    "def score(output_path: Path) -> Dict:\n",
    "    return run_metrics(\n",
    "        gt_data=gt_data, \n",
    "        pred_data=load_ctc_data(output_path),\n",
    "        matcher=CTCMatched,\n",
    "        metrics=[CTCMetrics],\n",
    "    )[\"CTCMetrics\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27ed1d89-5f3e-49ed-995f-f7ab3bc8e61d",
   "metadata": {},
   "source": [
    "## Parameter Search\n",
    "\n",
    "Here, we evaluate the segmentation and tracking given multiple values of `gamma`, used on the normalization step before the Cellpose prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "457eac7d-850c-417e-889e-991a0695af0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_labels = []\n",
    "metrics = []\n",
    "gammas = [0.1, 0.25, 0.5, 1]\n",
    "sigma = 5.0\n",
    "\n",
    "for gamma in gammas:\n",
    "\n",
    "    # cellpose prediction\n",
    "    cellpose_labels = np.zeros_like(image, dtype=np.int32)\n",
    "    array_apply(\n",
    "        image,\n",
    "        out_array=cellpose_labels,\n",
    "        func=predict,\n",
    "        gamma=gamma,\n",
    "    )\n",
    "    all_labels.append(cellpose_labels)\n",
    "    \n",
    "    name = f\"{dataset}_labels_{gamma}\"\n",
    "    if not COLAB:\n",
    "        viewer.add_labels(cellpose_labels, name=name, visible=False)\n",
    "\n",
    "    # cell tracking using `labels` parameter, it's the same as using `labels_to_edges`.\n",
    "    track(\n",
    "        config,\n",
    "        labels=cellpose_labels,\n",
    "        sigma=sigma,\n",
    "        overwrite=True\n",
    "    )\n",
    "\n",
    "    # exporting to CTC format\n",
    "    output_path = Path(name.upper()) / \"TRA\"\n",
    "    to_ctc(output_path, config, overwrite=True)\n",
    "\n",
    "    # computing tracking score\n",
    "    metric = score(output_path)\n",
    "    metric[\"gamma\"] = gamma\n",
    "    metrics.append(metric)\n",
    "\n",
    "print(metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc19bac8-44a8-4e7b-8ecf-bfa3780e8f1f",
   "metadata": {},
   "source": [
    "## Combined Contours and Detection\n",
    "\n",
    "The `labels_to_edges` combines multiple segmentation labels into a single detection and contour map.\n",
    "\n",
    "The detection map is the maximum value between the binary masks of each label.\n",
    "\n",
    "The contour map is the average contour map of the binary contours of each label."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6701e7a1-ddc7-4e69-aa3c-4721fd7a84dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "detection, contours = labels_to_edges(all_labels, sigma=sigma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eeb0e902-636b-4feb-9982-5f6c8595ec7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "if COLAB:\n",
    "    display(stackview.curtain(image, detection))\n",
    "else:\n",
    "    layer = viewer.add_labels(detection)\n",
    "    screenshot()\n",
    "    layer.visible = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ab81612-3cd2-4533-af90-043b66421982",
   "metadata": {},
   "outputs": [],
   "source": [
    "if COLAB:\n",
    "    display(stackview.curtain(image, contours))\n",
    "else:\n",
    "    layer = viewer.add_image(contours)\n",
    "    screenshot()\n",
    "    layer.visible = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d2efd71-4476-4027-a22e-7eb82d8762f8",
   "metadata": {},
   "source": [
    "## Tracking\n",
    "\n",
    "Run the tracking algorithm on the provided configuration, and combined detections and contours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06f11f38-1b8c-4744-a1a9-56144c75b32e",
   "metadata": {},
   "outputs": [],
   "source": [
    "track(\n",
    "   config,\n",
    "   detection=detection,\n",
    "   edges=contours,\n",
    "   overwrite=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "435f4b7f-11e0-4ae1-a95e-0443c642f03d",
   "metadata": {},
   "source": [
    "Compute metrics for the multiple hypotheses tracking and compare the scores of the different approaches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9841030b-48e4-481e-9c8f-7d5dfc342d3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_path = Path(f\"{dataset}_COMBINED\") / \"TRA\"\n",
    "to_ctc(output_path, config, overwrite=True)\n",
    "\n",
    "metric = score(output_path)\n",
    "metrics.append(metric)\n",
    "\n",
    "df = pd.DataFrame(metrics)\n",
    "df.to_csv(f\"{dataset}_scores.csv\", index=False)\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "932b5b22-20d3-409b-bf53-a9003db7c623",
   "metadata": {},
   "source": [
    "## Exporting and Visualization\n",
    "\n",
    "The intermediate tracking data is stored on disk and must be exported to your preferred format.\n",
    "Here we convert the resulting tracks to a DataFrame and Zarr to visualize using napari if running locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "178ec1aa-831f-40ea-b8a9-c8fdd571ac72",
   "metadata": {},
   "outputs": [],
   "source": [
    "tracks_df, graph = to_tracks_layer(config)\n",
    "tracks_df.to_csv(f\"{dataset}_tracks.csv\", index=False)\n",
    "\n",
    "segments = tracks_to_zarr(\n",
    "    config,\n",
    "    tracks_df,\n",
    "    overwrite=True,\n",
    ")\n",
    "\n",
    "if COLAB:\n",
    "    display(stackview.curtain(image, segments))\n",
    "else:\n",
    "    viewer.add_tracks(\n",
    "        tracks_df[[\"track_id\", \"t\", \"y\", \"x\"]],\n",
    "        name=\"tracks\",\n",
    "        graph=graph,\n",
    "        visible=True,\n",
    "    )\n",
    "\n",
    "    viewer.add_labels(segments, name=\"segments\").contour = 2\n",
    "    screenshot()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "jupytext": {
   "cell_metadata_filter": "-all",
   "encoding": "# -*- coding: utf-8 -*-",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}