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

# # Object tracking using Bayesian Tracker
#
# This notebook shows how to previously localised and classified objects from a series of mask images and unite them over time using Bayesian multi-object tracking (btrack). 
#
# The sections of this notebook are as follows:
#
# 1. Load segmentation objects and filter based on size
# 2. Run btrack, uniting the objects locations over time
# 3. Batch process
#
# The data used in this notebook is timelapse microscopy data with h2b-gfp/rfp markers that show the spatial extent of the nucleus and it's mitotic state. 

import btrack
from tqdm.auto import tqdm
import os
import glob

# ## 1. Load segmentation objects

expt = 'ND0011'
pos = 'Pos6'
root_dir = '/home/nathan/data'

with btrack.dataio.HDF5FileHandler(
     f'{root_dir}/{expt}/{pos}/objects.h5', 'r', obj_type='obj_type_1',
) as hdf:
    objects = hdf.objects

# Check objects feature classification

objects[0]

# Excluding segments that are too small to feasibly be cells

objects = [o for o in objects if o.properties['area']>100.]

# ## 2. Run btrack  
#
# Unite each object with it's subsequent position at the following time point and export as a tracking file

# initialise a tracker session using a context manager
with btrack.BayesianTracker() as tracker:

    # configure the tracker using a config file
    tracker.configure_from_file(
        './models/MDCK_config_wildtype.json'
    )
    tracker.max_search_radius = 40

    # append the objects to be tracked
    tracker.append(objects)

    # set the volume
    tracker.volume=((0, 1200), (0, 1600), (-1e5, 1e5))

    # track them (in interactive mode)
    tracker.track_interactive(step_size=100)

    # generate hypotheses and run the global optimizer
    tracker.optimize()

    tracker.export(f'{root_dir}/{expt}/{pos}/tracks.h5', obj_type='obj_type_1')

    # get the tracks in a format for napari visualization (optional)
    visaulise_tracks, properties, graph = tracker.to_napari(ndim=2)
    
    tracks = tracker.tracks

# # 3. Batch process
# Iterate over many experiments and positions (need to ensure you define normalisation and classification functions above first)

# +
root_dir = '/home/nathan/data'
expt_list = ['ND0009', 'ND0010', 'ND0011']
pos_list = 'all'
overwrite = False

for expt in tqdm(expt_list):
    
        # Find all positions in that experiment, if pos_list is all then it finds all positions
        if pos_list == 'all':
            pos_list = [pos for pos in os.listdir(f'{root_dir}/{expt}') 
                    if 'Pos' in pos 
                    and os.path.isdir(f'{root_dir}/{expt}/{pos}')]  
            
        ### Iterate over all positions in that experiment
        for pos in tqdm(pos_list):
            
            ### check if overwrite param is false check if raw directory already created and if type of transform file already exists and decide whether to skip pos
            if not overwrite and glob.glob(f'{root_dir}/{expt}/{pos}/*tracks*.h5'):
                print(glob.glob(f'{root_dir}/{expt}/{pos}/*tracks*.h5'), f'file found, skipping {expt}/{pos}')
                continue
                
            print(f'Starting {expt}/{pos}')
            
            # load objects
            with btrack.dataio.HDF5FileHandler(
                 f'{root_dir}/{expt}/{pos}/objects.h5', 'r', obj_type='obj_type_1',
            ) as hdf:
                objects = hdf.objects
            
            # filter objects for non-cell type sizes
            objects = [o for o in objects if o.properties['area']>100.]

            # initialise a tracker session using a context manager
            with btrack.BayesianTracker() as tracker:

                # configure the tracker using a config file
                tracker.configure_from_file(
                    './models/MDCK_config_wildtype.json'
                )
                tracker.max_search_radius = 40

                # append the objects to be tracked
                tracker.append(objects)

                # set the volume
                tracker.volume=((0, 1200), (0, 1600), (-1e5, 1e5))

                # track them (in interactive mode)
                tracker.track_interactive(step_size=100)

                # generate hypotheses and run the global optimizer
                tracker.optimize()

                tracker.export(f'{root_dir}/{expt}/{pos}/tracks.h5', obj_type='obj_type_1')
                
            print(f'Finished {expt}/{pos}')

StarcoderdataJupyterScriptsDedupFiltered