def test_tracklet_interpolate(real_tracklets):
data = np.stack(list(real_tracklets[0].values()))[:10]
inds = np.arange(len(data))
gap = 2
inds[len(inds) // 2:] += gap
tracklet = Tracklet(data, inds)
assert len(tracklet) == len(data)
new_tracklet = tracklet.interpolate(max_gap=1)
assert len(new_tracklet) == len(data)
new_tracklet = tracklet.interpolate(max_gap=gap)
assert len(new_tracklet) == len(data) + gap
missing_inds = list(set(range(inds.max())).difference(inds))
assert np.all(new_tracklet.data[missing_inds, :, 2] == 0.5)
def _reconstruct_tracks_as_tracklets(df):
"""
Parameters:
-----------
df: DataFrame
loaded from an .h5 tracks file (obtained from `stitch_tracklets()`)
"""
from deeplabcut.refine_training_dataset.stitch import Tracklet
tracklets = []
for _, group in df.groupby("individuals", axis=1):
temp = group.dropna(how="all")
inds = temp.index.to_numpy()
track = Tracklet(temp.to_numpy().reshape((len(temp), -1, 3)), inds)
track = track.interpolate(max_gap=len(group))
tracklets.append(track)
return tracklets