# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Viewing Tracks in Napari
# <NAME>, 12/01/2022
#
# Based on [Single Cell Tracking Napari Tutorial](https://napari.org/tutorials/tracking/cell_tracking.html). Data from [Cell Tracking Challenge](http://celltrackingchallenge.net/3d-datasets/)

# ### Packages

# +
import os
import napari

import numpy as np
import pandas as pd

from skimage.io import imread
from skimage.measure import regionprops_table
# -

# ### Definitions

PATH = 'Fluo-N3DH-CE/'
NUM_IMAGES = 195


# ### Define functions

def load_image(idx: int):
    """Load an image from the sequence.

    Parameters
    ----------
    idx : int
        Index of the image to load.

    Returns
    -------
    image : np.ndarray
       The image specified by the index, idx
    """
    filename = os.path.join(PATH, '01_GT/TRA', f'man_track{idx:0>3}.tif')
    return imread(filename)


def regionprops_plus_time(idx):
    """Return the unique track label, centroid and time for each track vertex.

    Parameters
    ----------
    idx : int
        Index of the image to calculate the centroids and track labels.

    Returns
    -------
    data_df : pd.DataFrame
       The dataframe of track data for one time step (specified by idx).
    """
    props = regionprops_table(stack[idx, ...], properties=('label', 'centroid'))
    props['frame'] = np.full(props['label'].shape, idx)
    return pd.DataFrame(props)


def root(node: int):
    """Recursive function to determine the root node of each subgraph.

    Parameters
    ----------
    node : int
        the track_id of the starting graph node.

    Returns
    -------
    root_id : int
       The track_id of the root of the track specified by node.
    """
    if full_graph[node] == 0:  # we found the root
        return node
    return root(full_graph[node])


# ### Get data to be displayed

# +
# Read in image data - manually annotated
stack = np.asarray([load_image(i) for i in range(NUM_IMAGES)])

# Find the centroid of each cell
data_df_raw = pd.concat(
    [regionprops_plus_time(idx) for idx in range(NUM_IMAGES)]
).reset_index(drop=True)
# sort the data lexicographically by track_id and time
data_df = data_df_raw.sort_values(['label', 'frame'], ignore_index=True)

# create the final data array: track_id, T, Z, Y, X
data = data_df.loc[
    :, ['label', 'frame', 'centroid-0', 'centroid-1', 'centroid-2']
].to_numpy()
# -

# ### View Manual Annotations

napari.view_image(stack, name='image')
napari.run()

# ### View Tracklets

napari.view_tracks(data, name='tracklets')
napari.run()

# ### Create graph to represent associations between tracks

lbep = np.loadtxt(os.path.join(PATH, '01_GT/TRA', 'man_track.txt'), dtype=np.uint)
full_graph = dict(lbep[:, [0, 3]])
graph = {k: v for k, v in full_graph.items() if v != 0}

# ### Get root node for lineage trees

roots = {k: root(k) for k in full_graph.keys()}
properties = {'root_id': [roots[idx] for idx in data[:, 0]]}

# ### Read in original image data

timelapse = np.asarray(
    [imread(os.path.join(PATH, '01', f't{i:0>3}.tif')) for i in range(NUM_IMAGES)]
)

# ### Visualise tracks and cells

# +
# scale factor for dimensions in TZYX order
SCALE = (1.0, 1.0, 0.09, 0.09)

viewer = napari.Viewer()
viewer.add_image(timelapse, scale=SCALE, name='Fluo-N3DH-CE')
viewer.add_tracks(data, properties=properties, graph=graph, scale=SCALE, name='tracks')
napari.run()
# -



StarcoderdataJupyterScriptsDedupFiltered