train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Identifying Cell Trajectories with MAGIK**\n",
    "\n",
    "Cell tracking is a crucial step in the analysis of time-lapse microscopy images. It allows us to study the behavior of individual cells over time, providing insights into cell division, migration, and other dynamic processes. \n",
    "\n",
    "This tutorial exemplifies how to train [MAGIK](https://www.nature.com/articles/s42256-022-00595-0) (Motion Analysis Through Graph Inductive Knowledge) on your own data for cell tracking.\n",
    "\n",
    "#### **1. Loading the Data**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Begin by defining the paths for the training images and their corresponding segmentation masks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_image_path = \"path/to/training/images\"\n",
    "train_segmentation_path = \"path/to/training/segmentation/masks\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the segmentation masks should be labeled images, with each pixel assigned a unique cell ID. Each cell should have a distinct identifier, while the background is labeled as 0."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the data using the code snippet below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import cv2\n",
    "\n",
    "def load_images(path):\n",
    "    images = []\n",
    "    for file in glob.glob(path + \"/*.tif\"): # (1)\n",
    "        image = cv2.imread(file, cv2.IMREAD_UNCHANGED) # (2)\n",
    "        images.append(image) # (3)\n",
    "\n",
    "    return images\n",
    "\n",
    "train_images = load_images(train_image_path)\n",
    "train_segmentations = load_images(train_segmentation_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This script \n",
    "**(1)** iterates over each file in the specified directory path containing a `.tif` extension (modify the extension if your images are in a different format), \n",
    "**(2)** reads the image using OpenCV's `cv2.imread()` function with the `IMREAD_UNCHANGED flag` indicating that the image should be loaded without any modification or conversion, \n",
    "and **(3)** appends the loaded image to the `images` list."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us visualize some frames from the training image sequence along with their corresponding segmentation mask.\n",
    "\n",
    "The script below plots a specified number of frames distributed evenly over the training image sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "number_of_frames = 5\n",
    "\n",
    "total_frames = len(train_segmentations) \n",
    "\n",
    "plot_interval = total_frames // number_of_frames \n",
    "frames_to_plot = [i for i in range(0, total_frames, plot_interval)]\n",
    "\n",
    "fig, axs = plt.subplots(2, number_of_frames + 1, figsize=(20, 6))\n",
    "fig.patch.set_facecolor('white')\n",
    "\n",
    "for i, frame in enumerate(frames_to_plot):\n",
    "    if i == 0:\n",
    "        axs[0, i].set_ylabel(\"Intensity image\", fontsize=16)\n",
    "        axs[1, i].set_ylabel(\"Segmentation\", fontsize=16)\n",
    "\n",
    "    axs[0, i].imshow(train_images[frame], cmap='gray')\n",
    "    axs[0, i].set_title(f\"Frame {frame}\", fontsize=16)\n",
    "    axs[0, i].tick_params(axis='both', which='both', bottom=False, top=False, left=False, right=False, labelleft=False, labelbottom=False)\n",
    "\n",
    "    # Plot segmentation\n",
    "    axs[1, i].imshow(train_segmentations[frame], cmap='tab20b')\n",
    "    axs[1, i].tick_params(axis='both', which='both', bottom=False, top=False, left=False, right=False, labelleft=False, labelbottom=False)\n",
    "\n",
    "plt.subplots_adjust(wspace=0.02, hspace=0.02)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **Creating a Graph From Segmented Images**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MAGIK models cell motion and interactions as a directed graph, where nodes represent segmented cells and edges connect spatially close cells across frames.\n",
    "\n",
    "We will implement the `GraphFromSegmentations` class to generate a graph from the segmented video frames:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from skimage import measure\n",
    "\n",
    "import torch\n",
    "from torch_geometric.data import Data\n",
    "\n",
    "\n",
    "class GraphFromSegmentations:\n",
    "    def __init__(self, connectivity_radius, max_frame_distance):\n",
    "        self.connectivity_radius = connectivity_radius # (1)\n",
    "        self.max_frame_distance = max_frame_distance # (2)\n",
    "\n",
    "    def __call__(self, images, segmentations, relation):\n",
    "        x, node_index_labels, frames = [], [], []\n",
    "        for frame, (image, segmentation) in enumerate(zip(images, segmentations)): # (3)\n",
    "            features, index_labels = self.compute_node_features(image, segmentation) # (4)\n",
    "\n",
    "            x.append(features) # (5)\n",
    "            node_index_labels.append(index_labels) # (6)\n",
    "            frames.append([frame] * len(features)) # (7)\n",
    "\n",
    "        x = np.concatenate(x)\n",
    "        node_index_labels = np.concatenate(node_index_labels)\n",
    "        frames = np.concatenate(frames)\n",
    "\n",
    "        edge_index, edge_attr = self.compute_connectivity(x, frames) # (8)\n",
    "        edge_ground_truth = self.compute_ground_truth( # (9)\n",
    "            node_index_labels, edge_index, relation\n",
    "        )\n",
    "\n",
    "        edge_index = edge_index.T\n",
    "        edge_attr = edge_attr[:, None]\n",
    "        edge_ground_truth = edge_ground_truth[:, None]\n",
    "\n",
    "        graph = Data( # (10)\n",
    "            x=torch.tensor(x, dtype=torch.float),\n",
    "            edge_index=torch.tensor(edge_index, dtype=torch.long),\n",
    "            edge_attr=torch.tensor(edge_attr, dtype=torch.float),\n",
    "            distance=torch.tensor(edge_attr, dtype=torch.float),\n",
    "            frames=torch.tensor(frames, dtype=torch.float),\n",
    "            y=torch.tensor(edge_ground_truth, dtype=torch.float),\n",
    "        )\n",
    "\n",
    "        return graph\n",
    "\n",
    "    def compute_node_features(self, image, segmentation):\n",
    "        labels = np.unique(segmentation)\n",
    "\n",
    "        x, indices = [], []\n",
    "        for label in labels[1:]:\n",
    "            mask = segmentation == label\n",
    "            props = measure.regionprops(mask.astype(np.int32), intensity_image=image)[0]\n",
    "\n",
    "            centroids = props.centroid / np.array(segmentation.shape)\n",
    "\n",
    "            x.append([*centroids])\n",
    "            indices.append(label)\n",
    "\n",
    "        return x, indices\n",
    "\n",
    "    def compute_connectivity(self, x, frames):\n",
    "        positions = x[:, :2]\n",
    "        distances = np.linalg.norm(positions[:, None] - positions, axis=-1)\n",
    "\n",
    "        frame_diff = (frames[:, None] - frames) * -1\n",
    "\n",
    "        mask = (distances < self.connectivity_radius) & ( \n",
    "            (frame_diff <= self.max_frame_distance) & (frame_diff > 0)\n",
    "        )\n",
    "\n",
    "        edge_index = np.argwhere(mask) \n",
    "        edge_attr = distances[mask] \n",
    "\n",
    "        return edge_index, edge_attr\n",
    "\n",
    "    def compute_ground_truth(self, indices, edge_index, relation):\n",
    "        sender = indices[edge_index[:, 0]] \n",
    "        receiver = indices[edge_index[:, 1]]\n",
    "        self_connections_mask = sender == receiver\n",
    "\n",
    "        if relation is None:\n",
    "            return self_connections_mask\n",
    "        \n",
    "        relation_indices = relation[:, [-1, 0]] \n",
    "        relation_indices = relation_indices[relation_indices[:, 0] != 0]\n",
    "\n",
    "        relation_mask = np.zeros(len(edge_index), dtype=bool)\n",
    "        for i, (s, r) in enumerate(zip(sender, receiver)):\n",
    "            if np.any((relation_indices == [s, r]).all(1)): \n",
    "                relation_mask[i] = True\n",
    "\n",
    "        return self_connections_mask | relation_mask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `GraphFromSegmentations` class is initialized with two parameters: **(1)** `connectivity_radius` and **(2)** `max_frame_distance`. These parameters play an important role in establishing the spatial and temporal thresholds necessary for determining connectivity between nodes within the graph structure.\n",
    "\n",
    "In the call method, `GraphFromSegmentations` receives two inputs: the segmented video frames (`segmentations`) and the parent-child relationships between cells (`relation`). **(3)** The method identifies separate objects in each frame of the segmented video data using their index labels. Next, **(4)** it calculates relevant features such as normalized centroids and eccentricity. **(5)** These features are stored in a set called `x`. The algorithm repeats this process for every object in the frame, creating a collection of node features (`x`), **(6)** index labels (`node_index_labels`), and **(7)** their corresponding frame numbers (`frames`), all poised for further processing.\n",
    "\n",
    "Leveraging the extracted node features, **(8)** `GraphFromSegmentations` proceeds to calculate pairwise distances between the positions of the nodes. Simultaneously, it computes the temporal difference between frames. Based on the specified thresholds (`connectivity_radius` and `max_frame_distance`), it identifies nodes that are both spatially and temporally close. The result is a set of edge indices (`edge_index`) and corresponding distances (`edge_attr`) representing the connectivity between nodes.\n",
    "\n",
    "Finally, **(9)** the ground-truth edges are computed. The generated graph includes a redundant number of edges with respect to the actual associations between cells. MAGIK aims to prune the redundant edges while retaining the true connections. Therefore, the ground truth for each edge is a binary value indicating whether an edge should connect two detections, i.e., an edge classification problem. `GraphFromSegmentations` defines the ground truth by comparing the node index labels and parent-child relationships. Firstly, it identifies self-connections where sender and receiver nodes have the same node index labels. Next, it explores the cell relationships to find relational connections, such as cell divisions. The ground truths are derived from the combination of self-connections and relational connections. \n",
    "\n",
    "**(10)** `GraphFromSegmentations` constructs a PyTorch Data object using node features, edge indices, attributes, distances, frames, and ground truth. This object encapsulates all necessary information for training and validationing."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instantiate the `GraphFromSegmentations` class with a connectivity radius of 0.2 (equivalent to 20% of the image size; this value may vary based on factors such as cell density and motion) and a maximum frame distance of 2, which allows for reconnecting cells that are not detected in consecutive frames."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_constructor = GraphFromSegmentations(connectivity_radius=0.2, max_frame_distance=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now construct the training graph using `graph_constructor`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_relation_path = \"path/to/training/relations.txt\"\n",
    "train_relation = np.loadtxt(train_relation_path, dtype=int) # (1)\n",
    "\n",
    "train_graph = graph_constructor(train_images, train_segmentations, train_relation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(1) Loads the relations.txt file, which defines parent-child relationships between cells. The file follows a comma-separated format:\n",
    "\n",
    "```\n",
    "cell_id, start_frame, end_frame, parent_id\n",
    "```\n",
    "\n",
    "A `parent_id` of 0 indicates that no parent is assigned. \n",
    "\n",
    "If your dataset includes parent-child relationships, provide a file with this format; otherwise, set the relation parameter to `None`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the following code to explore the graph data structure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Number of nodes:\", len(train_graph.x))\n",
    "print(\"Number of edges:\", len(train_graph.edge_index[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8, 8))\n",
    "\n",
    "for i, j in train_graph.edge_index.T:\n",
    "    plt.plot(\n",
    "        [train_graph.x[i, 1], train_graph.x[j, 1]],\n",
    "        [train_graph.x[i, 0], train_graph.x[j, 0]],\n",
    "        c=\"black\",\n",
    "        alpha=0.5,\n",
    "    )\n",
    "\n",
    "plt.scatter(\n",
    "    train_graph.x[:, 1],\n",
    "    train_graph.x[:, 0],\n",
    "    c=train_graph.frames,\n",
    "    cmap=\"viridis\",\n",
    "    zorder=10,\n",
    ")\n",
    "# label colorbar\n",
    "cb = plt.colorbar()\n",
    "cb.ax.set_title('Frame', fontsize=14)\n",
    "plt.xlabel(\"x\", fontsize=14)\n",
    "plt.ylabel(\"y\", fontsize=14)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This scatter plot depicts a graph with nodes represented as dots. The $x$ and $y$ coordinates represent the normalized node centroids. The color of each dot corresponds to the frame number, as shown on the color bar. The black lines on the plot illustrate the edges."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **Building a Dataset for Training**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "The training dataset consists of a single graph derived from the training video sequence. Although this may initially appear as limited data, it proves to be ample for effectively training the MAGIK model. To address the scarcity of data, we adopt a strategic approach of augmenting the training graph by splitting it into smaller temporal subgraphs.\n",
    "\n",
    "The `CellTracingDataset` implements this augmentation strategy by dividing the training graph into smaller subgraphs parameterized by the `window_size`parameter:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset\n",
    "\n",
    "\n",
    "class CellTracingDataset(Dataset):\n",
    "    def __init__(self, graph, window_size, dataset_size, transform=None, min_frame=0):\n",
    "        self.graph = graph\n",
    "\n",
    "        self.window_size = window_size # (1)\n",
    "        self.dataset_size = dataset_size\n",
    "\n",
    "        frames, edge_index = graph.frames, graph.edge_index\n",
    "        self.pair_frames = torch.stack(\n",
    "            [frames[edge_index[0, :]], frames[edge_index[1, :]]], axis=1\n",
    "        )\n",
    "        self.frames = frames\n",
    "        self.max_frame = frames.max()\n",
    "\n",
    "        self.transform = transform \n",
    "\n",
    "        self.min_frame = min_frame\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.dataset_size\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        frame_idx = np.random.randint(self.window_size + self.min_frame, self.max_frame + 1) # (2)\n",
    "\n",
    "        start_frame = frame_idx - self.window_size\n",
    "        node_mask = (self.frames >= start_frame) & (self.frames < frame_idx) # (3)\n",
    "        x = self.graph.x[node_mask] # (4)\n",
    "\n",
    "        edge_mask = (self.pair_frames >= start_frame) & (self.pair_frames < frame_idx) # (5)\n",
    "        edge_mask = edge_mask.all(axis=1) \n",
    "\n",
    "        edge_index = self.graph.edge_index[:, edge_mask] # (6)\n",
    "        edge_index -= edge_index.min() \n",
    "\n",
    "        edge_attr = self.graph.edge_attr[edge_mask] # (7)\n",
    "\n",
    "        # sample ground truth edges\n",
    "        ground_truth_edges = self.graph.y[edge_mask] # (8)\n",
    "\n",
    "        graph = Data( # (9)\n",
    "            x=x,\n",
    "            edge_index=edge_index,\n",
    "            edge_attr=edge_attr,\n",
    "            distance=edge_attr,\n",
    "            y=ground_truth_edges,\n",
    "        )\n",
    "\n",
    "        if self.transform: # (10)\n",
    "            graph = self.transform(graph)\n",
    "\n",
    "        return graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**(1)** The `window_size` parameter determines the number of frames in each subgraph. **(2)** The dataset generates subgraphs by randomly sliding a window across the training graph. The subgraph is constructed by extracting **(3-4)** nodes and **(5-8)** edges within the window. **(9)** The dataset returns the subgraph as a PyTorch Data object. \n",
    "\n",
    "To further enhance the training dataset, additional augmentations can be applied to the subgraphs. The `CellTracingDataset` class provides the flexibility to include custom augmentations by specifying the `transform` parameter. \n",
    "\n",
    "The following code snippet defines two augmentations: `RandomRotation` and `RandomFlip`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import math\n",
    "\n",
    "\n",
    "class RandomRotation: # (1)\n",
    "    def __call__(self, graph):\n",
    "        graph = graph.clone()\n",
    "        centered_features = graph.x[:, :2] - 0.5\n",
    "\n",
    "        angle = np.random.rand() * 2 * np.math.pi\n",
    "        rotation_matrix = torch.tensor(\n",
    "            [\n",
    "                [math.cos(angle), -math.sin(angle)],\n",
    "                [math.sin(angle), math.cos(angle)],\n",
    "            ]\n",
    "        )\n",
    "        rotated_features = torch.matmul(centered_features, rotation_matrix)\n",
    "\n",
    "        graph.x[:, :2] = rotated_features + 0.5\n",
    "        return graph\n",
    "    \n",
    "class RandomFlip: # (2)\n",
    "    def __call__(self, graph):\n",
    "        graph = graph.clone()\n",
    "        centered_features = graph.x[:, :2] - 0.5\n",
    "\n",
    "        if np.random.randint(2):\n",
    "            centered_features[:, 0] *= -1\n",
    "        \n",
    "        if np.random.randint(2):\n",
    "            centered_features[:, 1] *= -1\n",
    "        \n",
    "        graph.x[:, :2] = centered_features + 0.5\n",
    "        return graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**(1)** The `RandomRotation` augmentation function randomly rotates the positional features of the nodes within the subgraph. \n",
    "\n",
    "Likewise, **(2)** the `RandomFlip` augmentation randomly flips the positional features of the nodes along the $x$-axis or $y$-axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision import transforms\n",
    "\n",
    "train_dataset = CellTracingDataset(\n",
    "    train_graph,\n",
    "    window_size=5, # (1)\n",
    "    dataset_size=512, # (2)\n",
    "    transform=transforms.Compose([RandomRotation(), RandomFlip()]),\n",
    "    min_frame=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**(1)** The `window_size` parameter controls the number of frames in each subgraph generated from the training graph and fed into the model during training.\n",
    "**(2)** `dataset_size` controls the number of subgraphs generated from the training graph per epoch. \n",
    "**(3)** `min_frame` specifies the range of frames from which subgraphs are sampled. A value of 0 indicates that subgraphs can be extracted from the first frame through to the last frame of the video.\n",
    "\n",
    "The training dataset is instantiated with a window size of 5 frames and a dataset size of 512 subgraphs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **Defining the Data Loaders**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we proceed to define the data loaders, which are responsible for feeding the data to the model during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.data import DataLoader\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **Building MAGIK**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code snippet defines the [MAGIK](https://www.nature.com/articles/s42256-022-00595-0) model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from deeplay import GraphToEdgeMAGIK\n",
    "import torch.nn as nn\n",
    "\n",
    "model = GraphToEdgeMAGIK([96,] * 4, 1, out_activation=nn.Sigmoid) # (1)\n",
    "\n",
    "print(model) # (2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**(1)** Instantiates a simplified version of MAGIK, which is a message-passing neural network. The model has four layers, and each layer contains 96 hidden features. Along with the message-passing layers, the model also includes a node encoder, an edge encoder, and a classification head.  **(2)** Prints the model summary.\n",
    "\n",
    "MAGIK is similar to message-passing neural networks that we have seen in previous examples. However, the main difference is that MAGIK implements a local attention mechanism that allows the model to concentrate on specific nodes and edges during message passing. This mechanism comes into play when aggregating messages to a node. Each message's contribution has a weight that depends on the distance between the connected nodes through a function with learnable parameters defining a learnable local receptive field. With this mechanism, MAGIK can focus on relevant distance-based features during message passing, which is crucial for cell tracking tasks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **Training MAGIK**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the following code snippet to train or load a pre-trained MAGIK model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from deeplay import BinaryClassifier, Adam, Trainer\n",
    "\n",
    "classifier = BinaryClassifier(model=model, optimizer=Adam(lr=1e-3))\n",
    "classifier = classifier.create()\n",
    "\n",
    "trainer = Trainer(max_epochs=10)\n",
    "trainer.fit(classifier, train_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **Evaluating MAGIK**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the model is trained, we can evaluate its performance on the validation dataset. \n",
    "\n",
    "We start by loading the validation data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_image_path = \"path/to/val/images\"\n",
    "val_segmentation_path = \"path/to/val/segmentation/masks\"\n",
    "\n",
    "val_images = load_images(val_image_path)\n",
    "val_segmentations = load_images(val_segmentation_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and the corresponding parent-child relationships:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_relation_path = \"path/to/val/relations.txt\"\n",
    "val_relation = np.loadtxt(val_relation_path, dtype=int)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now construct the validation graph using `graph_constructor`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_graph = graph_constructor(val_images, val_segmentations, val_relation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Number of nodes:\", len(val_graph.x))\n",
    "print(\"Number of edges:\", len(val_graph.edge_index[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After creating the validation graph, we can assess the model's performance by calculating the f1-score of the predicted and ground-truth edge classification:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import f1_score\n",
    "\n",
    "classifier.eval()\n",
    "pred = classifier(val_graph)\n",
    "predictions = pred.detach().numpy() > 0.5\n",
    "\n",
    "ground_truth = val_graph.y\n",
    "\n",
    "score = f1_score(ground_truth, predictions)\n",
    "print(f\"Val F1 score: {score}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can expect an f1-score of approximately 0.99 on the validation graph, exhibiting the model's ability to accurately predict cell temporal associations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that MAGIK does not output cell trajectories, but a graph structure that shows the connections between cells across frames. To generate cell trajectories, a post-processing algorithm is applied to the predicted graph structure.\n",
    "\n",
    "The `compute_trajectories` function below implements a simple post-processing algorithm to compute cell trajectories from MAGIK predictions. This might be refined based on edge probability to improve the results, e.g., as done in the file `post.py` in the `\\SW\\lib` folder. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "import networkx as nx\n",
    "\n",
    "class compute_trajectories:\n",
    "\n",
    "    def __call__(self, graph, predictions):\n",
    "        pruned_edges = self.prune_edges(graph, predictions)\n",
    "\n",
    "        pruned_graph = nx.Graph()\n",
    "        pruned_graph.add_edges_from(pruned_edges)\n",
    "\n",
    "        trajectories = list(nx.connected_components(pruned_graph))\n",
    "\n",
    "        return trajectories\n",
    "\n",
    "    def prune_edges(self, graph, predictions):\n",
    "        pruned_edges = []\n",
    "\n",
    "        frame_pairs = np.stack(\n",
    "            [graph.frames[graph.edge_index[0]], graph.frames[graph.edge_index[1]]],\n",
    "            axis=1,\n",
    "        )\n",
    "\n",
    "        senders = np.unique(graph.edge_index[0])\n",
    "        for sender in senders: \n",
    "            sender_mask = graph.edge_index[0] == sender # (1)\n",
    "            candidate = predictions[sender_mask] == True # (2)\n",
    "\n",
    "            frame_diff = frame_pairs[sender_mask, 1] - frame_pairs[sender_mask, 0]\n",
    "            candidates_frame_diff = frame_diff[candidate]\n",
    "\n",
    "            if not np.any(candidate):\n",
    "                continue\n",
    "            else:\n",
    "                candidate_min_frame_diff = candidates_frame_diff.min()\n",
    "            \n",
    "            candidate_edge_index = graph.edge_index[:, sender_mask][ # (3)\n",
    "                :, candidate & (frame_diff == candidate_min_frame_diff)\n",
    "            ]\n",
    "            candidate_edge_index = candidate_edge_index.reshape(-1, 2)\n",
    "\n",
    "            if len(candidate_edge_index) == 1: # (4)\n",
    "                pruned_edges.append(tuple(*candidate_edge_index.numpy()))\n",
    "\n",
    "        return pruned_edges\n",
    "\n",
    "post_processor = compute_trajectories()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**(1)** The algorithm starts by selecting a node in the first frame ($t=0$) and then links it to other nodes in the following frames, **(2)** using only edges labeled as \"linked\" by MAGIK. **(3)** If there are no \"linked\" edges connecting the sender node at time $t$ to any receiver nodes at time $t+1$, the algorithm checks future frames up to a maximum time delay. If no \"linked\" edges are found within this timeframe, the trajectory ends.\n",
    "\n",
    "When a sender node has two \"linked\" edges connecting it to two receiver nodes in a later frame, it's identified as a division. In this case, **(4)** the algorithm creates two new trajectories. This process repeats until all \"linked\" edges are dealt with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "trajectories = post_processor(val_graph, predictions.squeeze())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we proceed to visualize the computed cell trajectories on top of the segmented video frames:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.animation import FuncAnimation\n",
    "from IPython.display import HTML\n",
    "from skimage import measure\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "size = val_segmentations[0].shape\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 8))\n",
    "\n",
    "list_of_colors = plt.cm.get_cmap(\"tab20b\", len(trajectories))\n",
    "np.random.shuffle(list_of_colors.colors)\n",
    "\n",
    "def update(frame):\n",
    "    ax.clear()\n",
    "    ax.imshow(val_images[frame], cmap=\"gray\")\n",
    "\n",
    "    segmentation = val_segmentations[frame]\n",
    "    labels = np.unique(segmentation)\n",
    "\n",
    "    for label in labels[1:]:\n",
    "        contour = measure.find_contours(segmentation == label, 0.5)[0]\n",
    "        ax.fill(\n",
    "            contour[:, 1],\n",
    "            contour[:, 0],\n",
    "            color=\"purple\",\n",
    "            alpha=0.2,\n",
    "            linewidth=6,\n",
    "        )\n",
    "\n",
    "    for idx, t in enumerate(trajectories):\n",
    "        coordinates = val_graph.x[list(t)]\n",
    "        frames = val_graph.frames[list(t)]\n",
    "\n",
    "        coordinates_in_frame = coordinates[frames == frame]\n",
    "\n",
    "        if len(coordinates_in_frame) == 0:\n",
    "            continue\n",
    "\n",
    "        ax.scatter(coordinates_in_frame[:, 1] * 512, coordinates_in_frame[:, 0] * 512, color=\"purple\")\n",
    "\n",
    "        coordinates_previous_frames = coordinates[\n",
    "                (frames <= frame) & (frames >= frame - 10)\n",
    "        ]\n",
    "        f = frames[(frames <= frame) & (frames >= frame - 10)]\n",
    "        coordinates_previous_frames = coordinates_previous_frames[np.argsort(f[f <= frame])]\n",
    "        ax.plot(\n",
    "            coordinates_previous_frames[:, 1] * 512,\n",
    "            coordinates_previous_frames[:, 0] * 512,\n",
    "            color=\"white\",\n",
    "        )\n",
    "\n",
    "        ax.plot(\n",
    "            coordinates_in_frame[max(0, frame - 10) : frame, 1] * 512,\n",
    "            coordinates_in_frame[max(0, frame - 10) : frame, 0] * 512,\n",
    "            color=\"red\",\n",
    "        )\n",
    "\n",
    "    return ax\n",
    "\n",
    "\n",
    "ani = FuncAnimation(fig, update, frames=len(val_segmentations))\n",
    "\n",
    "#html_video = HTML(ani.to_jshtml())\n",
    "html_video = HTML(ani.to_html5_video())\n",
    "\n",
    "plt.close()\n",
    "html_video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **Saving results**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The trained model can be saved for future use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.makedirs(\"model\", exist_ok=True)\n",
    "torch.save(classifier.model.state_dict(), \"model/model_state_dict.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These weights can be loaded as described in the \"test\" notebook, enabling you to evaluate the model on new data."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}