Using idTracker.ai with SLEAP

In [ ]:

!pip install sleap-io

!pip install sleap-io

Collecting sleap-io
  Downloading sleap_io-0.1.0-py3-none-any.whl (48 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0.0/48.5 kB ? eta -:--:--
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸━━━━━━ 41.0/48.5 kB 1.2 MB/s eta 0:00:01
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 48.5/48.5 kB 975.8 kB/s eta 0:00:00
Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from sleap-io) (1.25.2)
Requirement already satisfied: attrs in /usr/local/lib/python3.10/dist-packages (from sleap-io) (23.2.0)
Requirement already satisfied: h5py>=3.8.0 in /usr/local/lib/python3.10/dist-packages (from sleap-io) (3.9.0)
Collecting pynwb (from sleap-io)
  Downloading pynwb-2.6.0-py3-none-any.whl (135 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 135.4/135.4 kB 4.5 MB/s eta 0:00:00
Collecting ndx-pose (from sleap-io)
  Downloading ndx_pose-0.1.1-py2.py3-none-any.whl (8.3 kB)
Requirement already satisfied: pandas in /usr/local/lib/python3.10/dist-packages (from sleap-io) (2.0.3)
Collecting simplejson (from sleap-io)
  Downloading simplejson-3.19.2-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (137 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 137.9/137.9 kB 11.3 MB/s eta 0:00:00
Requirement already satisfied: imageio in /usr/local/lib/python3.10/dist-packages (from sleap-io) (2.31.6)
Requirement already satisfied: imageio-ffmpeg in /usr/local/lib/python3.10/dist-packages (from sleap-io) (0.4.9)
Collecting av (from sleap-io)
  Downloading av-12.0.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (33.8 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 33.8/33.8 MB 18.0 MB/s eta 0:00:00
Requirement already satisfied: pillow<10.1.0,>=8.3.2 in /usr/local/lib/python3.10/dist-packages (from imageio->sleap-io) (9.4.0)
Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from imageio-ffmpeg->sleap-io) (67.7.2)
Collecting hdmf<4,>=2.5.6 (from ndx-pose->sleap-io)
  Downloading hdmf-3.13.0-py3-none-any.whl (336 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 336.3/336.3 kB 21.9 MB/s eta 0:00:00
Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.10/dist-packages (from pynwb->sleap-io) (2.8.2)
Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas->sleap-io) (2023.4)
Requirement already satisfied: tzdata>=2022.1 in /usr/local/lib/python3.10/dist-packages (from pandas->sleap-io) (2024.1)
Requirement already satisfied: jsonschema>=2.6.0 in /usr/local/lib/python3.10/dist-packages (from hdmf<4,>=2.5.6->ndx-pose->sleap-io) (4.19.2)
Collecting ruamel-yaml>=0.16 (from hdmf<4,>=2.5.6->ndx-pose->sleap-io)
  Downloading ruamel.yaml-0.18.6-py3-none-any.whl (117 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 117.8/117.8 kB 501.6 kB/s eta 0:00:00
Requirement already satisfied: scipy>=1.4 in /usr/local/lib/python3.10/dist-packages (from hdmf<4,>=2.5.6->ndx-pose->sleap-io) (1.11.4)
Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.7.3->pynwb->sleap-io) (1.16.0)
Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=2.6.0->hdmf<4,>=2.5.6->ndx-pose->sleap-io) (2023.12.1)
Requirement already satisfied: referencing>=0.28.4 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=2.6.0->hdmf<4,>=2.5.6->ndx-pose->sleap-io) (0.34.0)
Requirement already satisfied: rpds-py>=0.7.1 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=2.6.0->hdmf<4,>=2.5.6->ndx-pose->sleap-io) (0.18.0)
Collecting ruamel.yaml.clib>=0.2.7 (from ruamel-yaml>=0.16->hdmf<4,>=2.5.6->ndx-pose->sleap-io)
  Downloading ruamel.yaml.clib-0.2.8-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_24_x86_64.whl (526 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 526.7/526.7 kB 32.2 MB/s eta 0:00:00
Installing collected packages: simplejson, ruamel.yaml.clib, av, ruamel-yaml, hdmf, pynwb, ndx-pose, sleap-io
Successfully installed av-12.0.0 hdmf-3.13.0 ndx-pose-0.1.1 pynwb-2.6.0 ruamel-yaml-0.18.6 ruamel.yaml.clib-0.2.8 simplejson-3.19.2 sleap-io-0.1.0

In [ ]:

!gdown --fuzzy https://drive.google.com/file/d/15iKXpcogB217kMXBKeqpcYXqVzGO4wdB/view?usp=sharing
!gdown --fuzzy https://drive.google.com/file/d/1k6SgvG_G41lhRRCK_LHaGEKLPBGS7fBQ/view?usp=drive_link

!gdown --fuzzy https://drive.google.com/file/d/15iKXpcogB217kMXBKeqpcYXqVzGO4wdB/view?usp=sharing !gdown --fuzzy https://drive.google.com/file/d/1k6SgvG_G41lhRRCK_LHaGEKLPBGS7fBQ/view?usp=drive_link

Downloading...
From (original): https://drive.google.com/uc?id=15iKXpcogB217kMXBKeqpcYXqVzGO4wdB
From (redirected): https://drive.google.com/uc?id=15iKXpcogB217kMXBKeqpcYXqVzGO4wdB&confirm=t&uuid=04cfd2c8-f81e-4ba6-88d9-19f048d389d4
To: /content/top-03132020125012-0000.mp4_1.slp
100% 32.0M/32.0M [00:00<00:00, 121MB/s] 
Downloading...
From: https://drive.google.com/uc?id=1k6SgvG_G41lhRRCK_LHaGEKLPBGS7fBQ
To: /content/without_gaps.npy
100% 2.52M/2.52M [00:00<00:00, 54.0MB/s]

In [ ]:

import numpy as np
import sleap_io as sio
from scipy.optimize import linear_sum_assignment

import numpy as np import sleap_io as sio from scipy.optimize import linear_sum_assignment

In [ ]:

labels = sio.load_file("top-03132020125012-0000.mp4_1.slp")
labels

labels = sio.load_file("top-03132020125012-0000.mp4_1.slp") labels

Out[ ]:

Labels(labeled_frames=52393, videos=1, skeletons=1, tracks=9)

In [ ]:

# Remove original tracks
for lf in labels:
    for inst in lf:
        inst.track = None
labels.tracks = []
labels

# Remove original tracks for lf in labels: for inst in lf: inst.track = None labels.tracks = [] labels

Out[ ]:

Labels(labeled_frames=52393, videos=1, skeletons=1, tracks=0)

In [ ]:

idt = np.load("without_gaps.npy", allow_pickle=True)
idt = idt[()]
idt

idt = np.load("without_gaps.npy", allow_pickle=True) idt = idt[()] idt

Out[ ]:

{'trajectories': array([[[         nan,          nan],
         [         nan,          nan]],
 
        [[902.3781965 , 387.54912517],
         [         nan,          nan]],
 
        [[908.45450435, 386.46998983],
         [         nan,          nan]],
 
        ...,
 
        [[562.99513437, 714.13227038],
         [         nan,          nan]],
 
        [[561.62233206, 713.61222197],
         [         nan,          nan]],
 
        [[562.02744577, 713.82905415],
         [         nan,          nan]]]),
 'version': '5.2.11',
 'video_paths': ['/Users/soline/Documents/idtrackerai/demo/top-03132020125012-0000.mp4'],
 'frames_per_second': 30,
 'body_length': 366.0,
 'stats': {'estimated_accuracy': 0.9918651663500414,
  'estimated_accuracy_after_interpolation': 0.9973305904537996,
  'percentage_identified': 0.3909587158589888,
  'estimated_accuracy_identified': 0.9973305904537996},
 'areas': {'mean': array([3932.50205202, 3597.3280052 ]),
  'median': array([3955.5, 3578.5]),
  'std': array([567.59482599, 547.1421946 ])},
 'setup_points': {},
 'identities_labels': ['1', '2'],
 'identities_groups': {},
 'length_unit': None,
 'id_probabilities': array([[[nan],
         [nan]],
 
        [[ 1.],
         [nan]],
 
        [[ 1.],
         [nan]],
 
        ...,
 
        [[ 1.],
         [nan]],
 
        [[ 1.],
         [nan]],
 
        [[ 1.],
         [nan]]])}

In [ ]:

idt_trx = idt["trajectories"]
idt_trx.shape  # (frames, tracks, xy)

idt_trx = idt["trajectories"] idt_trx.shape # (frames, tracks, xy)

Out[ ]:

(52393, 2, 2)

In [ ]:

# Create SLEAP tracks for each idtracker track.
n_idt_tracks = idt_trx.shape[1]
idt_tracks = [sio.Track(name=f"idtracker_{i}")for i in range(n_idt_tracks)]

# Save to labels.
labels.tracks = idt_tracks

# Create SLEAP tracks for each idtracker track. n_idt_tracks = idt_trx.shape[1] idt_tracks = [sio.Track(name=f"idtracker_{i}")for i in range(n_idt_tracks)] # Save to labels. labels.tracks = idt_tracks

In [ ]:

for lf in labels:

    if len(lf) == 0:
        continue

    # Get SLEAP centroids.
    pts = lf.numpy()  # (instances, nodes, xy)
    slp_centroids = np.nanmedian(pts, axis=1)  # (instances, xy)

    # Get idtracker centroids.
    idt_centroids = idt_trx[lf.frame_idx]  # (tracks, xy)
    not_missing = ~(np.isnan(idt_centroids).any(axis=-1))
    idt_track_inds = np.argwhere(not_missing).reshape(-1)  # (n_tracks,)
    idt_centroids = idt_centroids[not_missing]

    if len(idt_centroids) == 0:
        continue

    # Match centroids.
    cost = np.linalg.norm(
        np.expand_dims(idt_centroids, axis=1) -
        np.expand_dims(slp_centroids, axis=0),
        axis=-1
    )
    matches_idt, matches_slp = linear_sum_assignment(cost)

    # Assign idtracker tracks.
    for idt_ind, slp_ind in zip(matches_idt, matches_slp):
        lf[slp_ind].track = idt_tracks[idt_track_inds[idt_ind]]

for lf in labels: if len(lf) == 0: continue # Get SLEAP centroids. pts = lf.numpy() # (instances, nodes, xy) slp_centroids = np.nanmedian(pts, axis=1) # (instances, xy) # Get idtracker centroids. idt_centroids = idt_trx[lf.frame_idx] # (tracks, xy) not_missing = ~(np.isnan(idt_centroids).any(axis=-1)) idt_track_inds = np.argwhere(not_missing).reshape(-1) # (n_tracks,) idt_centroids = idt_centroids[not_missing] if len(idt_centroids) == 0: continue # Match centroids. cost = np.linalg.norm( np.expand_dims(idt_centroids, axis=1) - np.expand_dims(slp_centroids, axis=0), axis=-1 ) matches_idt, matches_slp = linear_sum_assignment(cost) # Assign idtracker tracks. for idt_ind, slp_ind in zip(matches_idt, matches_slp): lf[slp_ind].track = idt_tracks[idt_track_inds[idt_ind]]

In [ ]:

labels.save("idtracked.slp")

labels.save("idtracked.slp")