Converter script for storing trackmate outputs as .slp files

trackmate2slp.py

"""Converter script for storing trackmate outputs as .slp files

Usage:
`python trackmate2slp.py -l [PATH_TO_LABELS_1] [PATH_TO_LABELS_2] [PATH_TO_LABELS_N] -v [PATH_TO_TIF_1] [PATH_TO_TIF_2] [PATH_TO_TIF_3] -o [OUTDIR] -n`
"""
import pandas as pd
import numpy as np
import sleap
from skimage.io import imread
from pathlib import Path
import argparse
from tqdm import tqdm
def make_labels(video_path: str, trajectories_path: str) -> sleap.Labels:
    """Convert trackmate file to sleap Labels object

    Args:
      video_path: path to .tif video file
      trajectories_path: path to trackmate spots file to convert

    Returns:
      labels: the trackmate labels as a `sleap.Labels` object
    """
    vid = sleap.Video.from_numpy(imread(video_path))
    traj = pd.read_csv(trajectories_path)
    traj = traj.apply(pd.to_numeric, errors='coerce', downcast='integer')
    traj=traj.drop(range(0,3),axis=0)
    traj['TRACK_ID'].fillna(-1,inplace=True)
    print(sorted(traj["TRACK_ID"].unique()))
    skel = sleap.Skeleton()
    skel.add_node("centroid")
    tracks = {i:sleap.Track(0, name=f"Track {int(i)}") if i != -1 else sleap.Track(0, name=f"Unassigned Track") for i in sorted(traj["TRACK_ID"].unique())}
    lfs = []
    for frame_idx in sorted(traj["FRAME"].unique()):
        insts = []
        lf = traj[traj["FRAME"]==frame_idx]
        for inst_idx in sorted(lf['TRACK_ID'].unique()):
            instance = lf[lf['TRACK_ID']==inst_idx]
            pt = np.array((instance["POSITION_X"].iloc[0],instance["POSITION_Y"].iloc[0]))
            if np.isnan(pt).all():
                print("Nan found")
            try:
                insts.append(sleap.Instance.from_pointsarray(pt.reshape(1, 2), skeleton=skel, track=tracks[int(inst_idx)]))
            except Exception as e:
                print(inst_idx)
                raise(e)
        lfs.append(sleap.LabeledFrame(video=vid, frame_idx=frame_idx, instances=insts))
    labels = sleap.Labels(lfs)
    return labels
def tif2npy(video_path: str,save: bool=True) -> np.array:
    """Convert tif to numpy array

    Args:
      video_path: path to .tif file to convert
      save: A boolean flag indicating whether or not to save to .npy

    Returns:
      vid: video as an np array
    """
    vid = imread(video_path)
    vid = np.expand_dims(vid,-1)
    if save:
        np.save(f"{args.out_dir[-1]}/{Path(video_path).stem}.npy",vid)
    return vid

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--label_files','-l',action='store',nargs='+',help='path to trackmate csv files')
    parser.add_argument('--vid_files', '-v', action='store',nargs='+',help='path to tif files files')
    parser.add_argument('--to_npy','-n',action='store_true',help='Flag to convert tif to .npy file')
    parser.add_argument('--out_dir','-o',action='store',nargs=1,help='Where to store the .npy and .slp files',default='.')
    args = parser.parse_args()
    if len(args.label_files) !=0 and len(args.vid_files) != 0:
        for csv_file, tif in tqdm(zip(args.label_files,args.vid_files)):
            labels = make_labels(tif,csv_file)
            labels.save(f'{args.out_dir[-1]}/{Path(tif).stem}.slp')
    if args.to_npy:
        for tif in args.vid_files:
            tif2npy(tif)